package ext.sim.contact;

import ext.sim.Body;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.LinkedList;
import math.matSparse;
import static math.matSparse.*;
import math.math;
import math.vec3;
import static math.vec3.*;
import math.vecN;
import static math.vecN.*;

public class ContactSolver
{
    
    private static final double RESTITUTION = 0.6;
    
    private static final double FRICTION_COEFF = 0.3;
    
    private static final int PGS_ITER = 20;
    
    
    
    public ContactSolver()
    {
    }
    
    
    public void solve(
            final ArrayList<Body> aBodies, 
            final LinkedList<Contact> aContacts, 
            double aDT)
    {
        int K = aContacts.size();
        if (K == 0)
        {
            return;
        }
        int N = aBodies.size();
        
        matSparse Minv = new matSparse(6 * N, 6 * N);
        for (int n = 0; n < N; n++)
        {
            Body b = aBodies.get(n);
            
            Minv.set(6*n + 0, 6*n + 0, b.mIsPinned ? 0.0 : 1.0 / b.mMass);
            Minv.set(6*n + 1, 6*n + 1, b.mIsPinned ? 0.0 : 1.0 / b.mMass);
            Minv.set(6*n + 2, 6*n + 2, b.mIsPinned ? 0.0 : 1.0 / b.mMass);
            
            Minv.set(6*n + 3, 6*n + 3, b.mIsPinned ? 0.0 : b.mInvI.m[0]);
            Minv.set(6*n + 3, 6*n + 4, b.mIsPinned ? 0.0 : b.mInvI.m[3]);
            Minv.set(6*n + 3, 6*n + 5, b.mIsPinned ? 0.0 : b.mInvI.m[6]);
            
            Minv.set(6*n + 4, 6*n + 3, b.mIsPinned ? 0.0 : b.mInvI.m[1]);
            Minv.set(6*n + 4, 6*n + 4, b.mIsPinned ? 0.0 : b.mInvI.m[4]);
            Minv.set(6*n + 4, 6*n + 5, b.mIsPinned ? 0.0 : b.mInvI.m[7]);
            
            Minv.set(6*n + 5, 6*n + 3, b.mIsPinned ? 0.0 : b.mInvI.m[2]);
            Minv.set(6*n + 5, 6*n + 4, b.mIsPinned ? 0.0 : b.mInvI.m[5]);
            Minv.set(6*n + 5, 6*n + 5, b.mIsPinned ? 0.0 : b.mInvI.m[8]);
        }
        
        vecN impulseExt = new vecN(6 * N);
        for (int n = 0; n < N; n++)
        {
            Body b = aBodies.get(n);
            
            impulseExt.set(6*n + 0, b.mF.x * Minv.get(6*n + 0) * aDT); // LINEAR FORCE
            impulseExt.set(6*n + 1, b.mF.y * Minv.get(6*n + 1) * aDT);
            impulseExt.set(6*n + 2, b.mF.z * Minv.get(6*n + 2) * aDT);
            
            impulseExt.set(6*n + 3,                                    // ANGULAR FORCE
                (
                    Minv.get(6*n + 3, 6*n + 3) * b.mT.x + 
                    Minv.get(6*n + 3, 6*n + 4) * b.mT.y + 
                    Minv.get(6*n + 3, 6*n + 5) * b.mT.z
                ) * aDT
            );
            impulseExt.set(6*n + 4, 
                (
                    Minv.get(6*n + 4, 6*n + 3) * b.mT.x + 
                    Minv.get(6*n + 4, 6*n + 4) * b.mT.y + 
                    Minv.get(6*n + 4, 6*n + 5) * b.mT.z
                ) * aDT
            );
            impulseExt.set(6*n + 5, 
                (
                    Minv.get(6*n + 5, 6*n + 3) * b.mT.x + 
                    Minv.get(6*n + 5, 6*n + 4) * b.mT.y + 
                    Minv.get(6*n + 5, 6*n + 5) * b.mT.z
                ) * aDT
            );
        }
        
        vecN velocity = new vecN(6 * N);
        for (int n = 0; n < N; n++)
        {
            Body b = aBodies.get(n);
            
            velocity.set(6*n + 0, b.mV.x); // LINEAR VELOCITY
            velocity.set(6*n + 1, b.mV.y);
            velocity.set(6*n + 2, b.mV.z);
            
            velocity.set(6*n + 3, b.mW.x); // ANGULAR VELOCITY
            velocity.set(6*n + 4, b.mW.y);
            velocity.set(6*n + 5, b.mW.z);
        }
        
        matSparse J = new matSparse(3 * K, 6 * N);
        for (int k = 0; k < K; k++)
        {
            Contact c = aContacts.get(k);
            
            int i = c.mI;
            int j = c.mJ;
            
            J.set(k + 0, 6*i + 0, -c.mNormal.x); // FILL IN BLOCK FOR BODY I:
            J.set(k + 0, 6*i + 1, -c.mNormal.y);
            J.set(k + 0, 6*i + 2, -c.mNormal.z);
            vec3 RIxN = cross(c.mRi, c.mNormal);
            J.set(k + 0, 6*i + 3, -RIxN.x);
            J.set(k + 0, 6*i + 4, -RIxN.y);
            J.set(k + 0, 6*i + 5, -RIxN.z);
            
            J.set(k + 1, 6*i + 0, -c.mTangent.x);
            J.set(k + 1, 6*i + 1, -c.mTangent.y);
            J.set(k + 1, 6*i + 2, -c.mTangent.z);
            vec3 RIxT = cross(c.mRi, c.mTangent);
            J.set(k + 1, 6*i + 3, -RIxT.x);
            J.set(k + 1, 6*i + 4, -RIxT.y);
            J.set(k + 1, 6*i + 5, -RIxT.z);
            
            J.set(k + 2, 6*i + 0, -c.mBitangent.x);
            J.set(k + 2, 6*i + 1, -c.mBitangent.y);
            J.set(k + 2, 6*i + 2, -c.mBitangent.z);
            vec3 RIxB = cross(c.mRi, c.mBitangent);
            J.set(k + 2, 6*i + 3, -RIxB.x);
            J.set(k + 2, 6*i + 4, -RIxB.y);
            J.set(k + 2, 6*i + 5, -RIxB.z);
            
            
            J.set(k + 0, 6*j + 0, c.mNormal.x); // FILL IN BLOCK FOR BODY J:
            J.set(k + 0, 6*j + 1, c.mNormal.y);
            J.set(k + 0, 6*j + 2, c.mNormal.z);
            vec3 RJxN = cross(c.mRj, c.mNormal);
            J.set(k + 0, 6*j + 3, RJxN.x);
            J.set(k + 0, 6*j + 4, RJxN.y);
            J.set(k + 0, 6*j + 5, RJxN.z);
            
            J.set(k + 1, 6*j + 0, c.mTangent.x);
            J.set(k + 1, 6*j + 1, c.mTangent.y);
            J.set(k + 1, 6*j + 2, c.mTangent.z);
            vec3 RJxT = cross(c.mRj, c.mTangent);
            J.set(k + 1, 6*j + 3, RJxT.x);
            J.set(k + 1, 6*j + 4, RJxT.y);
            J.set(k + 1, 6*j + 5, RJxT.z);
            
            J.set(k + 2, 6*j + 0, c.mBitangent.x);
            J.set(k + 2, 6*j + 1, c.mBitangent.y);
            J.set(k + 2, 6*j + 2, c.mBitangent.z);
            vec3 RJxB = cross(c.mRj, c.mBitangent);
            J.set(k + 2, 6*j + 3, RJxB.x);
            J.set(k + 2, 6*j + 4, RJxB.y);
            J.set(k + 2, 6*j + 5, RJxB.z);
        }
        
        matSparse Jtranspose = J.copy();
        Jtranspose.transpose();
        
        matSparse Minv_Jtranspose = mul(Minv, Jtranspose);
        
        matSparse A = mul(J, Minv_Jtranspose);
        
        vecN Jv = mul(J, velocity); // Size = 3*K
        vecN b = add(Jv, mul(J, impulseExt)); // Size = 3*K
        for (int k = 0; k < K; k++)
        {
            b.set(3*k, b.get(3*k) + RESTITUTION * Jv.get(3*k));
        }
        
        vecN lambda = new vecN(3 * K);
        
        for (int iter = 0; iter < PGS_ITER; iter++)
        {
            for (int i = 0; i < lambda.m.length; i++)
            {
                double coeff = -1.0 / A.get(i, i);
                double sum = 0.0;
                HashMap<Integer, Double> ArowI = A.getRow(i);
                for (int j : ArowI.keySet())
                {
                    if (i == j)
                    {
                        continue;
                    }
                    double Aij = ArowI.get(j);
                    sum += Aij * lambda.get(j);
                }
                
                lambda.set(i, coeff * (sum + b.get(i)));
                
                if (i % 3 == 0) // Normal direction.
                {
                    lambda.set(i, math.clampd(lambda.get(i), 0.0, Double.MAX_VALUE));
                }
                else
                {
                    int iNormal = i / 3; // Use the normal direction lambda component.
                    double maxFric = FRICTION_COEFF * lambda.get(iNormal);
                    
                    lambda.set(i, math.clampd(lambda.get(i), -maxFric, maxFric));
                }
            }
        }
        
        
        for (int n = 0; n < N; n++)
        {
            Body body = aBodies.get(n);
            
            double impLinX = dotRowI_Vector(Minv_Jtranspose, 6*n + 0, lambda); // APPLY LINEAR IMPULSE.
            double impLinY = dotRowI_Vector(Minv_Jtranspose, 6*n + 1, lambda);
            double impLinZ = dotRowI_Vector(Minv_Jtranspose, 6*n + 2, lambda);
            body.mV.x += impLinX;
            body.mV.y += impLinY;
            body.mV.z += impLinZ;
            
            body.mP.set(body.mV); // Update Linear Momentum.
            body.mP.mul(body.mMass);
            
            double impAngX = dotRowI_Vector(Minv_Jtranspose, 6*n + 3, lambda); // APPLY ANGULAR IMPULSE.
            double impAngY = dotRowI_Vector(Minv_Jtranspose, 6*n + 4, lambda);
            double impAngZ = dotRowI_Vector(Minv_Jtranspose, 6*n + 5, lambda);
            body.mW.x += impAngX;
            body.mW.y += impAngY;
            body.mW.z += impAngZ;
            
            body.mL.set(body.mI.mul(body.mW)); // Update Angular Momentum.
        }
    }
    
    
    
    
}
